#  PYTHON 3.7

from pylab import *
import time
from matplotlib.backends.backend_pdf import PdfPages
import pandas as pd
import seaborn as sns
from pylab import *
import matplotlib.pyplot as plt
import os
import numpy as np

pd.set_option('display.expand_frame_repr', False)
pd.set_option('display.max_rows', 150)
pd.set_option('display.max_columns', 250)
pd.set_option('display.width', 320)

datadir = '/home/emd/Dropbox/Alem and Dugoua/ReplicationPackage/'
os.chdir(datadir)
myfile = pd.read_stata(datadir + '/sample_data.dta')
list(myfile.columns)
[x for x in list(myfile.columns) if 'q3' in x]
myfile = myfile.rename(columns={'Seed_gender': 'Seed Gender', 'treatment': 'Treatment', 'q302_WTPsol': 'WTP'})
myfile['Treatment'] = myfile['Treatment'].apply(lambda x: 'Unincentivized' if x == 'Network' else x)
myfile['Treatment'] = myfile['Treatment'].apply(lambda x: 'Incentivized' if x == 'Communication' else x)
myfile[['Seed Gender', 'Treatment', 'WTP']]


# SIMPLE HISTOGRAMS
# Histogram of bids for solar lanterns.
d = myfile['WTP']
sns.set()
sns.set(style="white", palette="muted", color_codes=True)
sns.distplot(d.dropna(), kde=False, rug=False, color="b")
plt.xlabel("Frequency")
plt.ylabel("Willingness to pay for solar lanterns (Rs)")
savefig(datadir + 'Graphs/histogram.pdf', bbox_inches='tight')
plt.close()


# FACETED HISTOGRAMS
# Faceted histogram of bids for solar lanterns across treatment groups
sns.set()
sns.set(style="whitegrid")
g = sns.FacetGrid(myfile, row="Seed Gender", col="Treatment", margin_titles=True)
bins = np.linspace(0, 1400, 15)
g = (g.map(plt.hist, "WTP", color="steelblue", bins=bins, lw=0))
g.set_xlabels("Willingness to pay (Rs)")
savefig(datadir + 'Graphs/facetedhistograms.pdf', bbox_inches='tight')
plt.close()


# SUBSIDIES
# Projected adoption rate of solar lanterns as a function of offered subsidy
cumfre = pd.DataFrame(columns={'All', 'Control', 'Unincentivized', 'Incentivized'})
maindata = myfile[['Treatment', 'WTP']]
data = maindata
# data['WTP'].describe()
data['Subsidy'] = 1200 - data['WTP']
# data['Subsidy'].describe()
crosstab = pd.crosstab(index=data["Subsidy"], columns="Count")
crosstab['Cumulative Count'] = crosstab['Count'].cumsum()
crosstab['Cumulative Freq'] = round(100 * crosstab['Cumulative Count'] / crosstab['Count'].sum(), 2)
cumfre['All'] = crosstab['Cumulative Freq']

for treat in ['Control', 'Unincentivized', 'Incentivized']:
    mask = myfile['Treatment'] == treat
    data = maindata[mask]
    # data['WTP'].describe()
    data['Subsidy'] = 1200 - data['WTP']
    # data['Subsidy'].describe()
    crosstab = pd.crosstab(index=data["Subsidy"], columns="Count")
    crosstab['Cumulative Count'] = crosstab['Count'].cumsum()
    crosstab['Cumulative Freq'] = round(100 * crosstab['Cumulative Count'] / crosstab['Count'].sum(), 2)
    cumfre[treat] = crosstab['Cumulative Freq']

cumfre = cumfre[['All', 'Control', 'Unincentivized', 'Incentivized']]
list_idx = list(cumfre.index)
cumfre = cumfre.fillna(0)
for index, row in cumfre.iterrows():
    print(index)
    for treat in ['Control', 'Unincentivized', 'Incentivized']:
        if cumfre.loc[index, treat] == 0:
            if list_idx.index(index) != 0:
                newindex = list_idx[list_idx.index(index)-1]
                cumfre.loc[index, treat] = cumfre.loc[newindex, treat]

del cumfre['All']
# import seaborn as sns
sns.set(style="whitegrid")
sns.lineplot(data=cumfre, palette="Blues", linewidth=2.5)
plt.xlabel("Subsidy (Rs)")
plt.ylabel("Projected adoption (%)")
savefig(datadir + 'Graphs/subsidy.pdf', bbox_inches='tight')
plt.close()


# Google trends
pal = sns.light_palette("navy", reverse=True)
# sns.palplot(pal)
# India weekly
gootren_UP = pd.read_csv(datadir + '2014-today_solartermsIndiaweekly.csv', sep=',', header=1)
gootren_UP['all'] = gootren_UP.iloc[:, 1] + gootren_UP.iloc[:, 2] + gootren_UP.iloc[:, 4]
gootren_UP['all'] = gootren_UP.iloc[:, 1] + gootren_UP.iloc[:, 2]
gootren_UP['yrmnth'] = gootren_UP['Week'].str[:7]
gootren_UP['Week'] = pd.to_datetime(gootren_UP['Week'])
gootren_UP = gootren_UP[gootren_UP['Week'] < pd.to_datetime('2016-01-01')]
gootren_UP[gootren_UP['yrmnth'] == '2015-10']
gootren_UP[gootren_UP['yrmnth'] == '2015-09']
gootren_UP[gootren_UP['yrmnth'] == '2015-07']
fig, ax = plt.subplots(figsize=(11, 6))
ax.plot(gootren_UP.iloc[:, 0], gootren_UP.iloc[:, 5], '-', label='solar lantern or lamp', linewidth=1, color=pal[0])
ax.plot(gootren_UP.iloc[:, 0], gootren_UP.iloc[:, 3], '-', label='solar lighting', linewidth=1, color=pal[3])
ax.axvline(x=pd.to_datetime('2015-07-26'), color='grey', linestyle='--')
ax.text(pd.to_datetime('2015-07-26'), 100, 'Beginning', rotation=90, fontsize=10, color='grey')
ax.axvline(x=pd.to_datetime('2015-10-04'), color='grey', linestyle='--')
ax.text(pd.to_datetime('2015-10-04'), 100, 'End', rotation=90, fontsize=10, color='grey')
ax.legend(ncol=1, loc="best", frameon=False)
ax.set(ylabel='Relative Frequency of Search in Google', xlabel='Week')
fig.savefig(datadir + 'Graphs/googletrends_india_week.png', bbox_inches='tight')
plt.close()


gootren_UP = pd.read_csv(datadir + 'comparedtokerosene.csv', sep=',', header=1)
gootren_UP['Week'] = pd.to_datetime(gootren_UP['Week'])
gootren_UP = gootren_UP[gootren_UP['Week'] < pd.to_datetime('2016-01-01')]
fig, ax = plt.subplots(figsize=(11, 6))
ax.plot(gootren_UP.iloc[:, 0], gootren_UP.iloc[:, 1], '-', label='solar lantern or lamp', linewidth=1, color=pal[0])
ax.plot(gootren_UP.iloc[:, 0], gootren_UP.iloc[:, 2], '-', label='kerosene', linewidth=1, color=pal[4])
ax.axvline(x=pd.to_datetime('2015-07-26'), color='grey', linestyle='--')
ax.text(pd.to_datetime('2015-07-26'), 80, 'Beginning', rotation=90, fontsize=10, color='grey')
ax.axvline(x=pd.to_datetime('2015-10-04'), color='grey', linestyle='--')
ax.text(pd.to_datetime('2015-10-04'), 80, 'End', rotation=90, fontsize=10, color='grey')
ax.legend(ncol=1, loc="best", frameon=False)
ax.set(ylabel='Relative Frequency of Search in Google', xlabel='Week')
fig.savefig(datadir + 'Graphs/googletrends_india_week_vskerosene.png', bbox_inches='tight')
plt.close()





#########################################################################################################
# APPENDIX
#########################################################################################################

# Box-plots of the amount of savings in rupees per treatment group
sns.set()
sns.set(style="ticks")
sns.boxplot(x="Treatment", y='q219_ifyes_howmuchinr', hue="Seed Gender", data=myfile, palette="Blues")
sns.despine(offset=10, trim=True)
plt.xlabel("Treatment groups")
plt.ylabel("Amount of savings (Rs)")
savefig(datadir + 'Graphs/boxplot_savings.pdf', bbox_inches='tight')
plt.close()


# Hexabin plot of WTP and amount of savings in Rupees.
# FOR VARIABLES: SH_sum_11_log and electrified_hh_nbr_log
# Hexbin plot with marginal distributions
myfile1 = myfile.rename(columns={'q302_yourfinalbidissolarlanterni': 'WTP', 'q219_ifyes_howmuchinr': 'Amount of Savings'})
gr = sns.jointplot(x='Amount of Savings', y='WTP', data=myfile1, kind='hex', gridsize=15)
gr = gr.plot_joint(sns.regplot, scatter=False, lowess=True, line_kws={"color": 'grey', "lw": 1})
gr.savefig(os.path.join(datadir + 'Graphs/saving_hexabin.png'), bbox_inches='tight')
close()



